from termcolor import colored
import boto3
import sys
from datetime import datetime
import json

author = {
    "name":"gl4ssesbo1",
    "twitter":"https://twitter.com/gl4ssesbo1",
    "github":"https://github.com/gl4ssesbo1",
    "blog":"https://www.pepperclipp.com/"
}

needs_creds = True

variables = {
    "SERVICE": {
        "value":"sts",
        "required":"true",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "ROLEARN": {
        "value":"",
        "required":"true",
        "description":"The ARN of the Role to get the creds from"
    },
    "ROLESESSIONNAME": {
        "value":"",
        "required":"true",
        "description":"A name of choice for the Role. No need to be the correct name."
    },
    "WEB-IDENTITY-TOKEN": {
        "value": "",
        "required": "true",
        "description": "The OAuth 2.0 access token or OpenID Connect ID token that is provided by the identity provider. Your application must get this token by authenticating the user who is using your application with a web identity provider."
    },
    "POLICY": {
        "value":"",
        "required":"false",
        "description":"The file path of an IAM policy in JSON format that you want to use as an inline session policy."
    },
    "POLICYARNS": {
        "value":"",
        "required":"false",
        "description":"The Amazon Resource Names (ARNs) of the IAM managed policies that you want to use as managed session policies. The policies must exist in the same account as the role. You can provide up to 10 managed policy ARNs."
    },
    "PROVIDER-ID": {
        "value":"",
        "required":"false",
        "description":"The fully qualified host component of the domain name of the identity provider. Specify this value only for OAuth 2.0 access tokens."
    },
    "DURATIONSECONDS": {
        "value":"3600",
        "required":"true",
        "description":"The time in seconds the token will last. Min is 900 seconds, Max is 12 hours. Default is 3600."
    }
}

description = "Returns a set of temporary security credentials for users who have been authenticated in a mobile or web application with a web identity provider. Example providers include Amazon Cognito, Login with Amazon, Facebook, Google, or any OpenID Connect-compatible identity provider."

aws_command = "assume-role-with-web-identity --role-arn <role-arn> --role-session-name <role-session-name> --web-identity-token <web-identity-token> --profile <profile>"

def exploit(profile, workspace):
    try:
        now = datetime.now()
        dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
        file = "{}_sts_assume_role_with_web_identity".format(dt_string)
        filename = "./workspaces/{}/{}".format(workspace, file)

        rolesessionname = variables['ROLESESSIONNAME']['value']
        rolearn = variables['ROLEARN']['value']
        durationseconds = int(variables['DURATIONSECONDS']['value'])

        policy = variables['POLICY']['value']
        providerid = variables['PROVIDER-ID']['value']
        tokencode = variables['WEB-IDENTITY-TOKEN']['value']

        arguments = {}

        arguments['RoleArn'] = rolearn
        arguments['RoleSessionName'] = rolesessionname
        arguments['DurationSeconds'] = durationseconds

        if not tokencode == "":
            arguments['TokenCode'] = tokencode

        if not policy == "":
            with open(policy, 'r') as pol:
                arguments['Policy'] = json.load(policy)
                pol.close()

        if not providerid == "":
            arguments['ProviderId'] = providerid

        if not variables['POLICYARNS']['value'] == "":
            policyarns = []
            for parn in (variables['POLICYARNS']['value']).split(","):
                arn = {}
                arn['arn'] = parn
                policyarns.append(arn)

            arguments['PolicyArns'] = policyarns

        response = profile.assume_role_with_web_identity(**arguments)
        
        full_response = {}
        for key,value in response.items():
            if not key == "ResponseMetadata":
                full_response[key] = value
        
        output = ""
        output += (colored("------------------------------------------------\n", "yellow", attrs=['bold']))
        output += ("{}: {}\n".format(colored("AccessKeyId", "yellow", attrs=['bold']), full_response['Credentials']['AccessKeyId']))
        output += (colored("------------------------------------------------\n", "yellow", attrs=['bold']))
        output += "\n"

        for key,value in full_response.items():
            if key == "PackedPolicySize":
                output += ("\t{}: {}\n".format(colored(key, "yellow", attrs=['bold']), colored(str(value), "red")))
            else:
                output += ("\t{}: \n".format(colored(key, "yellow", attrs=['bold'])))
                for k,v in value.items():
                    output += ("\t\t{}: {}\n".format(colored(k, "red", attrs=['bold']), colored(str(v), "blue")))
            
            output += "\n"
        print(output)
        
    except profile.exceptions.MalformedPolicyDocumentException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.PackedPolicyTooLargeException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.IDPRejectedClaimException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.IDPCommunicationErrorException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.InvalidIdentityTokenException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.RegionDisabledException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.ExpiredTokenException as e:
        print(colored("[*] {}".format(e), "red"))

    except:
        e = sys.exc_info()
        print(colored("[*] {}".format(e), "red"))